__version__ = "1.1.0"
import logging

import paramiko
from ruamel.yaml import safe_load

from logtools.logtools import logfile_hander, console_handler

logger = logging.getLogger(__name__)
logger.setLevel(level=logging.INFO)
# logger.root.setLevel(level=logging.INFO)
logger.addHandler(logfile_hander())
logger.addHandler(console_handler())


class SshTool:

    def __init__(self, ssh_ip, ssh_port=22, ssh_user='root', ssh_pwd='', key_path=None):
        """
        初始化
        :param ssh_ip:   IP
        :param ssh_port:  端口号
        :param ssh_user:  用户名
        :param ssh_pwd:  密码
        :param key_path: ssh 私钥
        """
        self.ip = ssh_ip
        self.port = ssh_port
        self.user = ssh_user
        self.pwd = ssh_pwd
        self.config = None
        if key_path is None:
            self.__load_config('../config/ssh_conf.yaml')
        else:
            self.__load_config(key_path)

    def __load_config(self, file_path):
        """
        加载全局配置文件
        :return:
        """
        with open(file_path, "r") as smallfile:
            config = safe_load(smallfile)  # pprint(config)
        self.config = config

    def get_ssh_connection(self):
        """
        使用用户名和密码登陆，仅推荐没有上传key的时候用
        :return:
        """
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        client.connect(hostname=self.ip, port=self.port, username=self.user, password=self.pwd, timeout=10)
        return client

    def get_key_connection(self):
        """
        推荐使用key登陆，很多命令不需要再次输入密码，方便
        :return:
        """
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        logger.debug('key_path: ', self.config['ssh_key'])
        private_key = paramiko.Ed25519Key.from_private_key_file(self.config['ssh_key'])
        client.connect(hostname=self.ip, username=self.user, pkey=private_key, timeout=10)
        return client

    def rsync_upload(self, local_dir='~/upload/*', upload_dir='/root/test/'):
        """
        在“服务器”上执行rsync命令，上传文件
        rsync -av ~/test/* root@ip:/root/test/
        :param local_dir: ~/upload/*
        :param upload_dir: /root/test/
        :return:
        """
        command = f"rsync -av {local_dir} {self.user}@{self.ip}:{upload_dir}"
        self.run_command(command)

    def rsync_download(self, server_dir='/root/test/*', local_dir='~/download/'):
        """
        在“服务器”上执行rsync命令，下载文件
        rsync -av root@ip:/root/test/* ~/test/
        :param server_dir: /root/test/*
        :param local_dir: ~/download/
        :return:
        """
        command = f"rsync -av {self.user}@{self.ip}:{server_dir} {local_dir}"
        self.run_command(command)

    def run_command(self, command="whoami"):
        """执行命令，执行时才建立连接，执行完退出连接。
        :param command:
        :return:如果输入多个命令，返回的是最后命令的状态。
        """
        conn = self.get_key_connection()
        logger.info("----------------------- 命令 ----------------------")
        logger.info(command)
        stdin, stdout, stderr = conn.exec_command(command=command)
        stdoutput = [line for line in stdout]
        stderroutput = [line for line in stderr]
        logger.info("----------------------- 结果 ----------------------")
        for output in stdoutput:
            logger.info(output.strip())
        for erroutput in stderroutput:
            logger.error(erroutput.strip())
        logger.info("Command status: %s" % (stdout.channel.recv_exit_status(),))
        if not stdout.channel.recv_exit_status():
            logger.debug("执行成功。")
            conn.close()
            if not stdoutput:
                stdoutput = True
            return True, stdoutput
        else:
            logger.error("执行失败！")
            for output in stderroutput:
                logger.error("%s" % (output,))
            conn.close()
            return False, stderroutput

    def run_sudo_command(self, command="whoami"):
        """执行sudo命令，执行时才建立连接，执行完退出连接。
        :param command:
        :return:如果输入多个命令，返回的是最后命令的状态。
        """
        conn = self.get_key_connection()
        logger.info("----------------------- 命令 ----------------------")
        logger.info(command)
        stdin, stdout, stderr = conn.exec_command(command=command)
        # 输入密码
        stdin.write(self.pwd + "\n")
        stdin.flush()
        stdoutput = [line for line in stdout]
        stderroutput = [line for line in stderr]
        logger.info("----------------------- 结果 ----------------------")
        for output in stdoutput:
            logger.info(output.strip())
        for erroutput in stderroutput:
            logger.error(erroutput.strip())
        logger.info("Command status: %s" % (stdout.channel.recv_exit_status(),))
        if not stdout.channel.recv_exit_status():
            logger.debug("执行成功。")
            conn.close()
            if not stdoutput:
                stdoutput = True
            return True, stdoutput
        else:
            logger.error("执行失败！")
            for output in stderroutput:
                logger.error("%s" % (output,))
            conn.close()
            return False, stderroutput

    def run_cmd_quit(self, command='whoami', quit_str='q'):
        """
        最后一个交互式命令，输入退出字符。多个交互式命令，不能同时用
        执行命令，执行时才建立连接，执行完退出连接。
        :param command:
        :param quit_str:  例如 q
        :return:如果输入多个命令，返回的是最后命令的状态。
        """
        conn = self.get_key_connection()
        logger.info("----------------------- 命令 ----------------------")
        logger.info(command)
        stdin, stdout, stderr = conn.exec_command(command=command)
        # 输入退出字符
        stdin.write(quit_str)
        stdin.flush()
        stdoutput = [line for line in stdout]
        stderroutput = [line for line in stderr]
        logger.info("----------------------- 结果 ----------------------")
        for output in stdoutput:
            logger.info(output.strip())
        for erroutput in stderroutput:
            logger.error(erroutput.strip())
        logger.info("Command status: %s" % (stdout.channel.recv_exit_status(),))
        if not stdout.channel.recv_exit_status():
            logger.debug("执行成功。")
            conn.close()
            if not stdoutput:
                stdoutput = True
            return True, stdoutput
        else:
            logger.error("执行失败！")
            for output in stderroutput:
                logger.error("%s" % (output,))
            conn.close()
            return False, stderroutput


if __name__ == "__main__":
    # ip = "172.25.140.101"
    # user = "root"
    # pwd = "Kml@0101"
    ip = "172.25.133.51"
    user = "root"
    pwd = "wentao"
    ssh = SshTool(ssh_ip=ip, ssh_user=user)
    # ssh = SshTool(ssh_ip=ip, ssh_user=user, key_path='../config/ssh_conf.yaml')
    ssh.run_command('id;pwd;')  # ssh.rsync_download()  # ssh.rsync_upload()

    # ssh.run_command('id;ll;')  # ssh.run_command('is;pwd;')  # ssh.run_command('is;pww;')

    # ssh.run_sudo_command('id;pwd;')  # ssh.run_sudo_command('id;ll;')  # ssh.run_sudo_command('is;pwd;')  # ssh.run_sudo_command('is;pww;')

    # ssh.run_cmd_quit('systemctl status docker','q')
